import backtrader as bt
import pandas as pd
import numpy as np
import datetime
import MySQLDataFeed as mdf
import SimpleMySqlClass as smc
import akshare as ak


class MomentumStrategy(bt.Strategy):
    params = (
        ("momentum_period", 126),  # 6个月动量周期 (交易日)
        ("holding_period", 21),  # 持有1个月
        ("rebalance_interval", 21),  # 每月调仓
        ("top_percentile", 0.1),  # 选择前10%的股票
        ("long_only", True),  # 仅做多
        ("printlog", True),
    )

    def __init__(self):
        # 初始化交易计数器
        self.order_counter = 0
        self.trade_dates = []

        # 为每个股票计算动量
        self.momentum = {}
        self.rankings = []
        for d in self.datas:
            self.momentum[d] = bt.indicators.Momentum(
                d.close, period=self.p.momentum_period
            )

    def next(self):
        # 每月第一个交易日执行调仓
        if len(self) % self.p.rebalance_interval != 0:
            return

        # 获取候选股票池
        available = [d for d in self.datas if len(d) > self.p.momentum_period]

        # 计算动量排名
        self.rankings = sorted(
            available, key=lambda d: self.momentum[d][0], reverse=True
        )

        # 选择前10%的股票
        top_size = int(len(self.rankings) * self.p.top_percentile)
        long_list = self.rankings[:top_size]

        # 平掉不在新名单中的持仓
        for d in self.getpositions():
            if d not in long_list and self.getposition(d).size > 0:
                self.close(d)

        # 等权重买入新标的
        portfolio_value = self.broker.getvalue()
        position_size = portfolio_value / len(long_list)

        for d in long_list:
            if self.getposition(d).size == 0:
                self.order_target_value(d, target=position_size)

        # 记录交易日期
        self.trade_dates.append(self.data.datetime.date())

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            return

        if order.status in [order.Completed]:
            if order.isbuy():
                log_type = "买入"
            elif order.issell():
                log_type = "卖出"

            self.log(
                f"{log_type}执行: {order.data._name} "
                f"数量: {order.executed.size} "
                f"价格: {order.executed.price:.2f} "
                f"成本: {order.executed.comm:.2f}"
            )

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log("订单取消/保证金不足/被拒绝")

    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            dt = dt or self.data.datetime.date()
            print(f"{dt.isoformat()}: {txt}")

    def stop(self):
        self.log(f"最终组合价值: {self.broker.getvalue():.2f}")


def get_stock_data(code, start_date, end_date):
    # stock_hfq_df = ak.stock_zh_a_hist(
    #     symbol=code, adjust="qfq", start_date=start_date, end_date=end_date
    # )
    stock_hfq_df = ak.stock_zh_a_hist(
        symbol="002202", adjust="qfq", start_date="20231201", end_date="20231215"
    )
    stock_hfq_df = stock_hfq_df.iloc[:, 0:6]
    stock_hfq_df.日期 = pd.to_datetime(stock_hfq_df.日期)
    print("1", stock_hfq_df.head())
    stock_hfq_df = stock_hfq_df.iloc[:, 0:6]
    print("2", stock_hfq_df.head())
    stock_hfq_df.日期 = pd.to_datetime(stock_hfq_df.日期)
    print("3", stock_hfq_df.head())

    # 日期     开盘     收盘     最高     最低     成交量  成交额    振幅   涨跌幅   涨跌额   换手率
    # 列名记得这样定义好
    stock_hfq_df.columns = ["Date", "Open", "Close", "High", "Low", "Volume"]
    print("4", stock_hfq_df.head())
    stock_hfq_df.set_index("Date", inplace=True)
    print("5", stock_hfq_df.head())
    return stock_hfq_df


def run_backtest():
    cerebro = bt.Cerebro()

    # 添加数据示例（需替换为真实数据）
    for code in ["600031", "600036", "600188", "600547"]:

        data = get_stock_data(code, "2024-01-01", "2024-12-31")
        cerebro.adddata(data)

    # 策略配置
    cerebro.addstrategy(MomentumStrategy)
    cerebro.broker.setcash(1000000)
    cerebro.broker.setcommission(commission=0.001)  # 0.1%佣金

    # 添加分析器
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name="sharpe")
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name="drawdown")

    # 运行回测
    results = cerebro.run()

    # 输出结果
    strat = results[0]
    print(f'夏普比率: {strat.analyzers.sharpe.get_analysis()["sharperatio"]:.2f}')
    print(
        f'最大回撤: {strat.analyzers.drawdown.get_analysis()["max"]["drawdown"]:.2f}%'
    )

    # 绘制结果
    cerebro.plot(style="candlestick")


if __name__ == "__main__":
    run_backtest()
